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

import argparse
import itertools
import logging
import sys
from datetime import datetime, timedelta, timezone

import sqlalchemy as sa
from sqlalchemy.exc import OperationalError, ProgrammingError

import lvestats.lib.commons.decorators
from lvestats.lib.bursting.history import HistoryShowBursting, TableDoesNotExistError
from lvestats.lib.commons.argparse_utils import ParseDatetime, check_non_negative_int
from lvestats.lib.commons.dateutil import gm_to_local, local_to_gm
from lvestats.lib.commons.func import get_ascii_table
from lvestats.lib.commons.logsetup import setup_logging
from lvestats.lib.config import ConfigError, read_config
from lvestats.lib.dbengine import make_db_engine, MakeDbException
from lvestats.orm import BurstingEventType, bursting_events_table
from lvestats.plugins.generic.burster.config import is_bursting_supported

log = setup_logging({},
                    caller_name='lve-bursting-info',
                    file_level=logging.WARNING,
                    console_level=logging.ERROR)


@lvestats.lib.commons.decorators.no_sigpipe
def main() -> None:
    if not is_bursting_supported():
        print('Bursting Limits feature is not supported in current environment')
        sys.exit(1)
    try:
        cfg = read_config()
        dbengine = make_db_engine(cfg=cfg)
        output = _main(cfg, dbengine, sys.argv[1:])
        if output:
            print(output)
    except TableDoesNotExistError as e:
        print(e.message)
        print('Enable Bursting Limits feature to create the table')
        sys.exit(1)
    except (OperationalError, ProgrammingError, ConfigError, MakeDbException) as e:
        log.error('Error occurred while executing the "lve-bursting-info" utility',
                  exc_info=e)
        sys.exit(1)


def _main(config: dict[str, str],
          dbengine: sa.engine.base.Engine,
          argv: list[str]) -> str:
    namespace = _parse_args(config, argv)

    # Convert date and time to UTC
    period_from_utc = local_to_gm(getattr(namespace, 'from'))  # 'from' is a reserved word
    period_to_utc = local_to_gm(namespace.to)

    history_show = HistoryShowBursting(
        dbengine=dbengine,
        period_from=period_from_utc,
        period_to=period_to_utc,
        uid=namespace.id,
        server_id=namespace.server_id,
    )
    rows = history_show.get()

    table = _build_table(rows, period_from_utc, period_to_utc)
    return table


def _parse_args(config: dict[str, str], argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Utility for displaying bursting limits statistics")

    parser.add_argument(
        '--id',
        type=check_non_negative_int,
        help='display records only for the specified LVE id')

    parser.add_argument(
        '--server-id',
        type=str,
        default=config.get('server_id', 'localhost'),
        help='used with central database for multiple servers, default "%(default)s"')

    # filtering output data by time period
    group_period = parser.add_argument_group()
    datetime_now = datetime.now()
    # default time delta is 1 hour
    # if the value is modified, don't forget to fix the CLI message
    default_timedelta = timedelta(hours=1)
    group_period.add_argument(
        '-f', '--from',
        action=ParseDatetime,
        default=datetime_now - default_timedelta,
        nargs='+',
        help='display records from the specified date and time; '
             'if not provided, will assume the last 1-hour time frame',
        metavar='YYYY-MM-DD[ HH:MM]')
    group_period.add_argument(
        '-t', '--to',
        action=ParseDatetime,
        default=datetime_now,
        nargs='+',
        help='display records up to the specified date and time; '
             'if not provided, will include records up to the current time',
        metavar='YYYY-MM-DD[ HH:MM]')

    return parser.parse_args(argv)


def _build_table(rows: list[sa.engine.RowProxy],
                 dt_from: datetime,
                 dt_to: datetime) -> str:
    """
    Build string representation of a console-displayable table
    using the rows retrieved from the DB table.
    """
    table_rows = _build_table_rows(rows,
                                   dt_from.timestamp(),
                                   dt_to.timestamp())
    table = get_ascii_table(table_rows,
                            fields=['ID', 'From', 'To'],
                            padding_width=1)
    return table


class TableRow(list):
    DEFAULT_FORMAT = '%Y-%m-%d %H:%M:%S.%f'

    def __init__(self,
                 lve_id: int,
                 start: float,
                 end: float,
                 format_: str = DEFAULT_FORMAT):
        # Convert date and time to local
        dt_start = gm_to_local(datetime.fromtimestamp(start, timezone.utc))
        dt_end = gm_to_local(datetime.fromtimestamp(end, timezone.utc))
        super().__init__([lve_id,
                          dt_start.strftime(format_),
                          dt_end.strftime(format_)])


def _build_table_rows(rows: list[sa.engine.RowProxy],
                      ts_from: float,
                      ts_to: float) -> list[TableRow]:
    """
    Build table rows representing time frame intervals
    when bursting was enabled for each LVE id.
    """
    result: list[TableRow] = []

    grouped_data = itertools.groupby(
        rows, key=lambda row: row[bursting_events_table.c.lve_id])

    for lve_id, group in grouped_data:
        group = list(group)

        first_event_idx = 0
        ts_from_, ts_to_ = None, None

        # If there is an event whose timestamp is less than 'ts_from',
        # attempting to detect whether bursting was already enabled
        # when the required time frame started
        if group[0][bursting_events_table.c.timestamp] < ts_from:
            first_event_idx = 1
            if group[0][bursting_events_table.c.event_type] == BurstingEventType.STARTED.value:
                ts_from_ = ts_from

        # Construct intervals for the certain LVE id
        for i in range(first_event_idx, len(group)):
            if group[i][bursting_events_table.c.event_type] == BurstingEventType.STARTED.value:
                if ts_from_ is None:
                    ts_from_ = group[i][bursting_events_table.c.timestamp]
                continue

            if ts_from_ is None:
                continue

            ts_to_ = group[i][bursting_events_table.c.timestamp]

            result.append(TableRow(lve_id, ts_from_, ts_to_))
            ts_from_, ts_to_ = None, None

        # Assuming bursting-enabled interval ends after the required time frame
        if ts_from_ is not None:
            result.append(TableRow(lve_id, ts_from_, ts_to))

    return result