Mini Shell

# coding=utf-8
#
# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2023 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENSE.TXT

# pylint: disable=abstract-method

import enum
from enum import IntEnum

import sqlalchemy as sa

from lvestats.orm.base import Base
from lvestats.orm.const import LVE_STATS_2_TABLENAME_PREFIX, SERVER_ID_LENGTH


class _EnumAsInteger(sa.types.TypeDecorator):
    # NOTE(vlebedev): Taken as-is from https://stackoverflow.com/a/38786737/3344105
    """Column type for storing Python enums in a database INTEGER column.

    This will behave erratically if a database value does not correspond to
    a known enum value.
    """

    impl = sa.types.SmallInteger  # underlying database type

    def __init__(self, enum_type):
        super().__init__()
        self.enum_type = enum_type

    def process_bind_param(self, value, dialect):
        if isinstance(value, self.enum_type):
            return value.value
        raise ValueError(
            f'expected {self.enum_type.__name__} value, got {value.__class__.__name__}'
        )

    def process_result_value(self, value, dialect):
        return self.enum_type(value)

    def copy(self, **kwargs):
        return _EnumAsInteger(self.enum_type)


@enum.unique
class BurstingEventType(IntEnum):
    STOPPED = 0
    STARTED = 1


bursting_events_table = sa.Table(
    LVE_STATS_2_TABLENAME_PREFIX + 'bursting_events',
    Base.metadata,
    # 'server_id' is essential for proper functionality
    # when used with a central database for multiple servers
    sa.Column('server_id', sa.String(SERVER_ID_LENGTH), primary_key=True),
    sa.Column('lve_id', sa.BigInteger, primary_key=True),
    sa.Column(
        'timestamp',
        sa.Integer,
        primary_key=True,
        index=True,
        comment='UNIX timestamp of the event (seconds precision)'
    ),
    sa.Column('event_type', _EnumAsInteger(BurstingEventType), nullable=False),
)